from machine import UART
import sensor, image, time 
from pyb import LED
import math
import tf

class UartCommend:
    def __init__(self, uart_port=1, baudrate=115200):
        """初始化 UART，设置帧头和帧尾，波特率等"""
        self.frame_header = 'header'  # 帧头
        self.frame_tail = 'tail'    # 帧尾
        self.uart = UART(uart_port, baudrate)  # 初始化 UART
        self.uart.init(baudrate, bits=8, parity=None, stop=1)  # 配置 UART 参数

    def send_uart_data(self, data, frame_header=None, frame_tail=None):
        """发送数据到 UART，并加上帧头和帧尾"""
        # 使用默认帧头和帧尾，如果提供了新的帧头和帧尾则使用它们
        if frame_header is None:
            frame_header = self.frame_header
        if frame_tail is None:
            frame_tail = self.frame_tail
        
        # 确保 data 是 bytearray 类型，并发送
        data = frame_header + data + frame_tail
        self.uart.write(data)

    def get_uart_data(self, frame_header=None, frame_tail=None):
        """从 UART 接收数据，并去掉帧头和帧尾"""
        # 使用默认帧头和帧尾，如果提供了新的帧头和帧尾则使用它们
        if frame_header is None:
            frame_header = self.frame_header
        if frame_tail is None:
            frame_tail = self.frame_tail
        
        # 读取串口数据
        data = self.uart.read()
        if data:
            # 去掉帧头和帧尾
            data = data.decode('utf-8')
            if data.startswith(frame_header) and data.endswith(frame_tail):
                data = data[len(frame_header):-len(frame_tail)]
                return data
        return None

class Camera:
    def __init__(self, roi=(80, 80, 120, 120), led_pin=4,pixformat=sensor.RGB565, framesize=sensor.QVGA):
        """初始化相机，设置分辨率、颜色格式、ROI区域和LED"""
        self.roi = roi  # 设ROI区域 (x, y, w, h)
        self.led = LED(led_pin)
        self.sensor = sensor  # 使用传入的传感器对象
        # 初始化相机
        self.sensor.reset()  # 重置并初始化相机
        self.sensor.set_pixformat(pixformat)  # 设置图像格式
        self.sensor.set_framesize(framesize)  # 设置图像分辨率
        self.sensor.skip_frames(time=2000)  # 等待设置生效

    def get_image(self):
        """获取并返回一帧图像，裁剪ROI区域并绘制矩形框"""
        img = self.sensor.snapshot()  # 捕获一帧图像
        
        # 绘制ROI区域的矩形框
        img.draw_rectangle(self.roi, color=(255, 0, 0), thickness=1)
        
        # 从图像中裁剪出ROI区域
        img_roi = img.crop(self.roi)

        # 返回裁剪后的图像和FPS值
        return img_roi


class ModelRunner:
    def __init__(self, model_path="/sd/model.tflite", label_path="/sd/label.txt"):
        self.model_path = model_path
        self.label_path = label_path
        self.tf_model = self.load_model()
        self.labels = self.load_labels()
        self.electronics = {'keyboard', 'mouse', 'headphones', 'monitor', 'speaker', 'printer', 'mobile_phone'}
        self.tools = {'wrench', 'soldering_iron', 'electrodrill', 'tape_measure', 'screwdriver', 'pliers', 'oscillograph', 'multimeter'}
        self.default_result = ('unknown', 'unknown')

    def load_model(self):
        print(f"Loading model: {self.model_path}")
        return tf.load(self.model_path)

    def load_labels(self):
        print(f"Loading labels: {self.label_path}")
        with open(self.label_path, "r") as f:
            return [line.strip() for line in f]

    def classify_image(self, image):
        results = []
        max_attempts = 10

        for _ in range(max_attempts):
            for obj in tf.classify(self.tf_model, image, min_scale=1.0, scale_mul=0.5, x_overlap=-1, y_overlap=-1):
                sorted_list = sorted(zip(self.labels, obj.output()), key=lambda x: x[1], reverse=True)
                top_result = sorted_list[0]

                label = top_result[0]
                confidence = top_result[1]

                if confidence > 0.7:
                    if label in self.electronics:
                        category = 'electronics'
                    elif label in self.tools:
                        category = 'tools'
                    else:
                        category = 'unknown'

                    results.append((label, category))

                if len(results) == max_attempts:
                    break

        if not results:
            return self.default_result

        small_class_counts = {}
        label_to_category = {}
        for label, category in results:
            if label not in small_class_counts:
                small_class_counts[label] = 0
            small_class_counts[label] += 1
            label_to_category[label] = category

        most_common_small_class = max(small_class_counts, key=small_class_counts.get)
        most_common_category = label_to_category[most_common_small_class]

        return (most_common_small_class, most_common_category)


